function df = SpectralGrad(varargin)

f    = varargin{1};
kind = varargin{2};
if length(varargin) == 3
    dim = varargin{3};
else
    dim = 1;
end

if strcmp(kind,'cheb')
    
    df = chebgrad(f,dim);

elseif strcmp(kind,'cheb2')
    
    df = chebgrad2(f,dim);

elseif strcmp(kind,'fourier')
    
    df = fouriergrad(f,dim);
    
else
   
    error('unknown kind') 
end

function dg = chebgrad(varargin)

% g = g(:,:,...1:N+1,:,:...);
%              dim
%assumes x = cos( pi*(2*n+1)/(2*N+2) )

g = varargin{1};

if numel(g) == max(size(g)) %1d derivatives
    dim = 1;
    g = g(:); %vectorize
else
    dim = varargin{2};
end

A = size(g);
M = A(dim); %M = number of elements in dim 

G = shiftdim(g,dim-1); %makes dim --> 1
K = size(G); K(1) = 2*K(1);
F = zeros(K);
 
F(1:M,:) = G(1:M,:); 
F(M+1:2*M,:) = G(M:-1:1,:);

Fhat = ifft(F,[],1);
Ftmp = zeros(size(G));
Ftmp(:) = Fhat(1:M,:);

k = 0:M-1; k = k(:);
k = k(:,ones(1,numel(g)/M)); %expand k so there are enough elements
k = reshape(k,size(G)); %reshape to appropriate size 

C = 2*exp(1i*pi*k/(2*M)).*Ftmp; C(1,:) = 1/2*C(1,:); %these are the cheby polynomial coefficients 

Fhat = zeros(K); 
Fhat(1:M,:) = k(1:M,:).*C(1:M,:).*exp(-1i*pi*k(1:M,:)/(2*M)); Fhat(1,:) = 2*Fhat(1,:); %pics out sin function in transform
Fhat(M+2:2*M,:) = -k(M:-1:2,:).*C(M:-1:2,:).*exp(1i*pi*k(M:-1:2,:)/(2*M)); %pics out sin function in transform

Ftmp = fft(Fhat,[],1)/(-2*1i); %factor of -2*1i comes from e^(-1i*p) - e^(1i*p) = -2*1i*sin(p)
x = cos(pi*(2*k+1)/(2*M));
F = zeros(size(G)); F(1:M,:) = Ftmp(1:M,:)./sqrt(1-x(1:M,:).^2);

if isreal(g) %only output real part if g is real
    dg = shiftdim(real(F),ndims(g) +1 - dim); %shift back to orginal dimensions
else
    dg = shiftdim(F,ndims(g) +1 - dim); %shift back to orginal dimensions
end


function dg = chebgrad2(varargin)

% g = g(:,:,...1:N+1,:,:...);
%              dim
%assumes x = cos( pi*(2*n+1)/(2*N+2) )

g = varargin{1};

if numel(g) == max(size(g)) %1d derivatives
    dim = 1;
    g = g(:); %vectorize
else
    dim = varargin{2};
end

A = size(g);
M = A(dim); %M = number of elements in dim 

G = shiftdim(g,dim-1); %makes dim --> 1
K = size(G); K(1) = 2*K(1)-2;
F = zeros(K);

F(1:M,:) = G(1:M,:); 
F(M+1:2*M -2,:) = G(M-1:-1:2,:);

Fhat = ifft(F,[],1);
Ftmp = zeros(size(G));
Ftmp(:) = Fhat(1:M,:);

k = 0:M-1; k = k(:);
k = k(:,ones(1,numel(g)/M)); %expand k so there are enough elements
k = reshape(k,size(G)); %reshape to appropriate size 

C = 2*Ftmp; C(1,:) = 1/2*C(1,:); C(end,:) = 1/2*C(end,:);%these are the cheby polynomial coefficients 

Fhat = zeros(K); 
%Fhat(1:M,:) = k(1:M,:).*C(1:M,:).*exp(-1i*pi*k(1:M,:)/(2*M)); Fhat(1,:) = 2*Fhat(1,:); %pics out sin function in transform
%Fhat(M+2:2*M,:) = -k(M:-1:2,:).*C(M:-1:2,:).*exp(1i*pi*k(M:-1:2,:)/(2*M)); %pics out sin function in transform

Fhat(1:M,:) = k(1:M,:).*C(1:M,:); Fhat(1,:) = 2*Fhat(1,:); %pics out sin function in transform
Fhat(M+1:2*M-2,:) = -k(M-1:-1:2,:).*C(M-1:-1:2,:); %pics out sin function in transform

Ftmp = fft(Fhat,[],1)/(-2*1i); %factor of -2*1i comes from e^(-1i*p) - e^(1i*p) = -2*1i*sin(p)
x = cos(pi*k/(M-1)); 
F = zeros(size(G)); F(1:M,:) = Ftmp(1:M,:)./sqrt(1-x(1:M,:).^2);

%%% fix nan in above formula %%%
tmp = sum(k.^2.*C,1);
F(1,:) = tmp(:);
tmp = sum(-((-1).^k).*k.^2.*C,1);
F(end,:) = tmp(:);

if isreal(g) %only output real part if g is real
    dg = shiftdim(real(F),ndims(g) +1 - dim); %shift back to orginal dimensions
else
    dg = shiftdim(F,ndims(g) +1 - dim); %shift back to orginal dimensions
end

function dg = fouriergrad(varargin)

% g = g(:,:,...1:N+1,:,:...);
%              dim
%assumes x = -pi + 2*pi*n/N, n = 0,1,...N-1

g = varargin{1};

if numel(g) == max(size(g)) %1d derivatives
    dim = 1;
    g = g(:); %vectorize
else
    dim = varargin{2};
end

A = size(g);
N = A(dim); %N = number of elements in dim 
M = (N - 1)/2;

F = shiftdim(g,dim-1); %makes dim --> 1

n = 0:N-1; n = n(:);
n = n(:,ones(1,numel(g)/N)); %expand k so there are enough elements
n = reshape(n,size(F));

k = n - M;

Fhat = exp(1i*pi*k).*fft(exp(2*pi*1i*n*M/N).*F,[],1); %these are the fourier coefficients 

F = exp(-2*pi*1i*n*M/N).*ifft(1i*k.*Fhat.*exp(-1i*k*pi),[],1); 

if isreal(g) %only output real part if g is real
    dg = shiftdim(real(F),ndims(g) +1 - dim); %shift back to orginal dimensions
else
    dg = shiftdim(F,ndims(g) +1 - dim); %shift back to orginal dimensions
end
